package cn.jbolt.ai.klb.service.impl;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.app.service.AiAppConfigService;
import cn.jbolt.ai.core.vdb.VdbCollectionService;
import cn.jbolt.ai.klb.entity.AiDoc;
import cn.jbolt.ai.klb.service.AiDocService;
import cn.jbolt.ai.klb.service.AiFileService;
import cn.jbolt.ai.resource.vdb.entity.VDBResource;
import cn.jbolt.ai.resource.vdb.service.VDBResourceService;
import cn.jbolt.util.Pager;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import cn.jbolt.ai.klb.dao.AiKnowledgeBaseMapper;
import cn.jbolt.ai.klb.entity.AiKnowledgeBase;
import cn.jbolt.ai.klb.service.AiKnowledgeBaseService;
import cn.jbolt.util.Result;
import com.jboltai.capability.vdb.JBoltVDB;
import com.jboltai.capability.vdb.VDBOptResult;
import com.jboltai.capability.vdb.collection.field.VDBField;
import com.jboltai.capability.vdb.collection.field.VDBFieldType;
import com.jboltai.capability.vdb.doc.VDBDoc;
import com.jboltai.capability.vdb.search.VDBConditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

@Service
public class AiKnowledgeBaseServiceImpl extends ServiceImpl<AiKnowledgeBaseMapper, AiKnowledgeBase> implements AiKnowledgeBaseService {
    private static final Logger logger = LoggerFactory.getLogger(AiKnowledgeBaseServiceImpl.class);

    @Autowired
    private VdbCollectionService vdbCollectionService;

    @Autowired
    private AiFileService aiFileService;

    @Autowired
    private AiDocService aiDocService;
    @Autowired
    private AiAppConfigService aiConfigService;
    @Autowired
    private VDBResourceService vdbResourceService;

    @Override
    public Result<AiKnowledgeBase> getKnowledgeBaseById(Long id) {
        AiKnowledgeBase knowledgeBase = this.getById(id);
        if (knowledgeBase != null) {
            return Result.success(knowledgeBase);
        } else {
            return Result.error("知识库不存在");
        }
    }

    @Override
    public Result listByPage(Boolean enable, String keywords, Integer pageNumber, Integer pageSize) {
        QueryWrapper<AiKnowledgeBase> queryWrapper = new QueryWrapper<>();
        if (enable != null) {
            queryWrapper.eq("enable", enable);
        }
        if (keywords != null && !keywords.isEmpty()) {
            queryWrapper.and(wrapper -> wrapper.like("name", keywords).
                    or().
                    like("brief_info", keywords));
        }
        queryWrapper.orderByDesc("id");
        return Result.success(Pager.from(this.page(new Page<>(pageNumber, pageSize), queryWrapper)));
    }

    @Override
    public Result listAllEnabled() {
        return Result.success(this.list(new QueryWrapper<AiKnowledgeBase>().eq("enable", true).orderByDesc("id")));
    }

    @Override
    @Transactional
    public Result add(AiKnowledgeBase knowledgeBase) {
        if (knowledgeBase.getVdbResourceId() == null) {
            return Result.error("请选择VDB资源库");
        }
        //查询vdb资源
        Result<VDBResource> result = vdbResourceService.getResourceById(knowledgeBase.getVdbResourceId());
        if (result.isError()) {
            return result;
        }
        VDBResource vdbResource = result.getData();
        knowledgeBase.setVdbResourceType(vdbResource.getType());
        knowledgeBase.setEnable(true);
        if (this.save(knowledgeBase)) {
            //创建完知识库以后,在生成对应的向量数据库
            Result vdbResult = createVDB(knowledgeBase.getId(), knowledgeBase.getVdbResourceId());
            if (vdbResult.isError()) {
                return vdbResult;
            }
            return Result.success();
        } else {
            return Result.error("添加失败");
        }
    }

    private Result createVDB(Long knowledgeBaseId, Long vdbResourceId) {
        //先检测集合是否已创建
        VDBOptResult<Boolean> res = JBoltVDB.use(vdbResourceId.toString()).checkCollectionExist("KDB_" + knowledgeBaseId.toString());
        if (res.isSuccess() == false) {
            return Result.error(res.getFullMsg());
        }
        if (res.getResult() == true) {
            return Result.success();
        }

        List<VDBField> vdbFields = Arrays.asList(
                VDBField.ofNormal("file_id", VDBFieldType.STRING),
                VDBField.ofNormal("data_type", VDBFieldType.STRING),
                VDBField.ofHnswVector(),
                VDBField.ofDoc(true)
        );
        //创建完知识库以后,在生成对应的向量数据库
        return vdbCollectionService.createCollection(vdbResourceId.toString(),
                "KDB_" + knowledgeBaseId.toString(), vdbFields);
    }


    @Override
    public Result update(AiKnowledgeBase knowledgeBase) {
        AiKnowledgeBase knowledgeBaseDB = getById(knowledgeBase.getId());
        if (!knowledgeBaseDB.getVdbResourceId().equals(knowledgeBase.getVdbResourceId())) {
            return Result.error("知识库中的VDB资源库不可以更改!");
        }
        boolean updated = updateById(knowledgeBase);
        if (updated) {
            Result vdbResult = createVDB(knowledgeBase.getId(), knowledgeBase.getVdbResourceId());
            if (vdbResult.isError()) {
                return vdbResult;
            }
            return Result.success();
        } else {
            return Result.error("更新失败");
        }
    }

    @Override
    @Transactional
    public Result del(Long id) {
        AiKnowledgeBase knowledgeBase = getById(id);
        if (knowledgeBase == null) {
            return Result.error("数据异常");
        }
        if (removeById(id)) {
            //移除相关的config ,更新知识库的数量
            Result configResult = aiConfigService.delConfigByKlbId(id);
            if (configResult.isError()) {
                return configResult;
            }
            //移除相关的file
            Result fileResult = aiFileService.delFileByKlbId(id);
            if (fileResult.isError()) {
                return fileResult;
            }
            //移除相关的doc
            Result docResult = aiDocService.delDocByKlbId(id);
            if (docResult.isError()) {
                return docResult;
            }
            //移除vdb中的表
            Result result = vdbCollectionService.removeCollection(knowledgeBase.getVdbResourceId().toString(), "KDB_" + knowledgeBase.getId());
            if (result.isError()) {
                return result;
            } else {
                return Result.success();
            }
        } else {
            return Result.error("删除失败");
        }
    }

    @Override
    public Result toggleEnable(Long id) {
        AiKnowledgeBase knowledgeBase = this.getById(id);
        if (knowledgeBase == null) {
            return Result.error("知识库不存在");
        }

        knowledgeBase.setEnable(!knowledgeBase.getEnable());
        boolean update = updateById(knowledgeBase);
        if (update) {
            Result vdbResult = createVDB(knowledgeBase.getId(), knowledgeBase.getVdbResourceId());
            if (vdbResult.isError()) {
                return vdbResult;
            }
            return Result.success();
        } else {
            return Result.error("更新启用状态失败");
        }
    }

    /**
     * 根据知识库id返回知识库
     *
     * @param knowledgeBaseId
     * @return
     */
    public Result<AiKnowledgeBase> checkKnowledgeBase(Long knowledgeBaseId) {
        if (StrUtil.isBlankIfStr(knowledgeBaseId)) {
            return Result.error("请选择知识库");
        }
        Result result = getKnowledgeBaseById(knowledgeBaseId);
        if (result.isError()) {
            return result;
        }
        AiKnowledgeBase aiKnowledgeBase = (AiKnowledgeBase) result.getData();

        if (!aiKnowledgeBase.getEnable()) {
            return Result.error(aiKnowledgeBase.getName() + "--知识库已关闭!");
        }
        return Result.success(aiKnowledgeBase);
    }

    public Result getListByVDBResourceId(Long vdbResourceId) {
        return Result.success(this.list(new QueryWrapper<AiKnowledgeBase>().eq("vdb_resource_id", vdbResourceId)));
    }



    @Override
    public Result<List<AiDoc>> getDocListByKnowledgeBaseId(Long knowledgeBaseId, String content, Integer limit, Double minScore) {
        Result<AiKnowledgeBase> result = checkKnowledgeBase(knowledgeBaseId);
        if (result.isError()) return Result.error(result.getCode(), result.getMsg());
        AiKnowledgeBase knowledgeBase = result.getData();
        try {
            VDBOptResult<List<VDBDoc>> VDBDocList = JBoltVDB.use(knowledgeBase.getVdbResourceId().toString()).
                    search("KDB_" + knowledgeBaseId, new VDBConditions().matchText(content).limit(limit));
            if (VDBDocList.isSuccess()) {
                List<VDBDoc> docsList = VDBDocList.getResult();
                if (docsList.size() == 0) {
                    return Result.success(new ArrayList<>());
                }
                //把docsList的中的id变成一个list
                List<String> idList = docsList.stream()
                        .filter(doc -> {
                            if (minScore == null ) {
                                return true;
                            }
                            return  ObjectUtil.defaultIfNull(doc.getScore(),0.0) >= minScore;
                        })
                        .map(VDBDoc::getId) // 假设 id 是 String 类型
                        .collect(Collectors.toList());
                List<AiDoc> aiDocsList = aiDocService.getDocListByIds(idList);

                syncScore(aiDocsList,docsList);

                return Result.success(aiDocsList);
            } else {
                return Result.error(VDBDocList.getMsg());
            }
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("知识库检索异常：" + e.getMessage());
        }
        return Result.error("知识库检索异常");
    }

    /**
     * 将 vdbDoc中的score同步到aiDocs中
     * @param aiDocs AI文档列表
     * @param vdbDoc 向量数据库文档列表
     */
    private void syncScore(List<AiDoc> aiDocs, List<VDBDoc> vdbDoc) {
        if (ObjectUtil.isEmpty(aiDocs) || ObjectUtil.isEmpty(vdbDoc)) {
            return;
        }
        Map<String, Double> vdbScoreMap = vdbDoc.stream()
            .collect(Collectors.toMap(
                VDBDoc::getId,
                VDBDoc::getScore,
                (existing, replacement) -> existing
            ));
        aiDocs.forEach(aiDoc -> {
            Double score = vdbScoreMap.get(aiDoc.getId().toString());
            if (score != null) {
                aiDoc.setScore(score);
            }
        });
    }
}
